package com.hyperboost.deepseek.store;

import com.hyperboost.deepseek.bean.ChatMessages;
import dev.langchain4j.data.message.ChatMessage;
import dev.langchain4j.data.message.ChatMessageDeserializer;
import dev.langchain4j.data.message.ChatMessageSerializer;
import dev.langchain4j.store.memory.chat.ChatMemoryStore;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.data.mongodb.core.MongoTemplate;
import org.springframework.data.mongodb.core.query.Criteria;
import org.springframework.data.mongodb.core.query.Query;
import org.springframework.data.mongodb.core.query.Update;
import org.springframework.stereotype.Component;

import java.util.LinkedList;
import java.util.List;

@Component
public class DeepSeekStore implements ChatMemoryStore {

    @Autowired
    private MongoTemplate mongoTemplate;

    /**
     * 获取聊天记忆
     *
     * @param memoryId 聊天记忆id
     * @return
     */
    @Override
    public List<ChatMessage> getMessages(Object memoryId) {
        Criteria criteria = Criteria.where("memoryId").is(memoryId);
        Query query = new Query(criteria);
        ChatMessages chatMessages = mongoTemplate.findOne(query, ChatMessages.class);
        // 判断是否查询到数据
        if (chatMessages == null) {
            return new LinkedList<>();
        }
        // 反序列化聊天消息列表，json -> list
//        String contentJson = chatMessages.getContent();
//        return ChatMessageDeserializer.messagesFromJson(contentJson);
//

        // 反序列化聊天消息列表，json -> list
        List<ChatMessage> chatMessageList = ChatMessageDeserializer.messagesFromJson(chatMessages.getContent());
        return chatMessageList;

    }

    /**
     * 更新聊天记忆
     *
     * @param memoryId 聊天记忆id
     * @param list     聊天消息列表
     */
    @Override
    public void updateMessages(Object memoryId, List<ChatMessage> list) {
        Criteria criteria = Criteria.where("memoryId").is(memoryId);
        Query query = new Query(criteria);
        Update update = new Update();

        // 序列化聊天消息列表，list -> json
        update.set("content", ChatMessageSerializer.messagesToJson(list));

        // 修改/新增数据：当查询不到数据时，则新增数据
        mongoTemplate.upsert(query, update, ChatMessages.class);
    }

    /**
     * 根据聊天记忆id删除聊天记忆
     *
     * @param memoryId 聊天记忆id
     */
    @Override
    public void deleteMessages(Object memoryId) {
        Criteria criteria = Criteria.where("memoryId").is(memoryId);
        Query query = new Query(criteria);
        mongoTemplate.remove(query, ChatMessages.class);
    }
}
